{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9492edb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from vit_pytorch import ViT\n",
    "    \n",
    "v = ViT(\n",
    "    image_size = 256,\n",
    "    channels = 1,\n",
    "    patch_size = 32,\n",
    "    num_classes = 5,\n",
    "    dim = 1024,\n",
    "    depth = 6,\n",
    "    heads = 16,\n",
    "    mlp_dim = 2048,\n",
    "    dropout = 0.1,\n",
    "    emb_dropout = 0.1\n",
    ")\n",
    "\n",
    "img = torch.randn(1, 1, 256, 256)\n",
    "preds = v(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6ecbca1b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.1757, -0.2250,  0.5245, -0.4812,  0.3510]],\n",
       "       grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "aabac809",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ViT(\n",
       "  (to_patch_embedding): Sequential(\n",
       "    (0): Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=32, p2=32)\n",
       "    (1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (2): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "    (3): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (dropout): Dropout(p=0.1, inplace=False)\n",
       "  (transformer): Transformer(\n",
       "    (norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (layers): ModuleList(\n",
       "      (0-5): 6 x ModuleList(\n",
       "        (0): Attention(\n",
       "          (norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (attend): Softmax(dim=-1)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "          (to_qkv): Linear(in_features=1024, out_features=3072, bias=False)\n",
       "          (to_out): Sequential(\n",
       "            (0): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (1): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (1): FeedForward(\n",
       "          (net): Sequential(\n",
       "            (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (1): Linear(in_features=1024, out_features=2048, bias=True)\n",
       "            (2): GELU(approximate='none')\n",
       "            (3): Dropout(p=0.1, inplace=False)\n",
       "            (4): Linear(in_features=2048, out_features=1024, bias=True)\n",
       "            (5): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (to_latent): Identity()\n",
       "  (mlp_head): Linear(in_features=1024, out_features=5, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a3e7f9f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "# import collections\n",
    "# import random\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision import models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8524abc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def replace_bn_with_identity(module:nn.Module):\n",
    "    # 遍历当前模块的所有子模块\n",
    "    for name, child in module.named_children(): \n",
    "        # 如果是 BN 层，则替换为 Identity \n",
    "        if isinstance(child, nn.BatchNorm2d):\n",
    "            setattr(module, name, nn.Identity())\n",
    "        else:\n",
    "            # 对非 BN 子模块递归调用\n",
    "            replace_bn_with_identity(child)\n",
    "\n",
    "def initialize_weights(module:nn.modules):\n",
    "    # 遍历当前模块的所有子模块\n",
    "        for name, child in module.named_children(): \n",
    "            # 如果是 BN 层，则替换为 Identity \n",
    "            if isinstance(child, nn.Conv2d) or isinstance(child, nn.Linear):\n",
    "                # nn.init.orthogonal_(module.weight, nn.init.calculate_gain('relu'))\n",
    "                nn.init.xavier_uniform_(child.weight)\n",
    "                # nn.init.kaiming_uniform_(module.weight)\n",
    "                if child.bias is not None:\n",
    "                    nn.init.constant_(child.bias, 0)\n",
    "            else:\n",
    "                replace_bn_with_identity(child)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "95b97fe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RESNET50_FC(nn.Module):\n",
    "    def __init__(self, input_channels=1, act_num=5, use_softmax=True):\n",
    "        super().__init__()\n",
    "        self.input_shape = (input_channels, 256, 256)\n",
    "        self.resnet = models.resnet50(weights=None) \n",
    "        replace_bn_with_identity(self.resnet)\n",
    "        self.resnet.conv1  = nn.Conv2d(input_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)\n",
    "        # self.resnet.avgpool = nn.MaxPool2d(5, stride=1)\n",
    "        self.resnet.fc  = nn.Identity()  # 使用 Identity 替代 fc 层\n",
    "        self.use_softmax = use_softmax\n",
    "        self.cov_out = nn.Sequential(\n",
    "            nn.Linear(2048, 1024),\n",
    "        )\n",
    "        self.critic_linear = nn.Linear(1024, 1)\n",
    "        self.actor_linear = nn.Linear(1024, act_num)\n",
    "        initialize_weights(self)\n",
    "\n",
    "    def forward(self, x:torch.Tensor):\n",
    "        for k, shape in enumerate(self.input_shape):\n",
    "            if x.shape[k+1] != shape:\n",
    "                raise ValueError(f\"Input shape should be {self.input_shape}, got {x.shape[1:]}\")\n",
    "        x = self.resnet(x)\n",
    "        x = F.relu(self.cov_out(x))\n",
    "        if self.use_softmax:\n",
    "            a = F.softmax(self.actor_linear(x), dim=1)\n",
    "        else:\n",
    "            a = self.actor_linear(x)\n",
    "        return a, self.critic_linear(x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c122c972",
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet = models.resnet50(weights=None) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "ad9d33cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "replace_bn_with_identity(resnet)\n",
    "resnet.conv1  = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)\n",
    "# resnet.avgpool = nn.AdaptiveAvgPool2d((2,2))\n",
    "resnet.fc  = nn.Identity()  # 使用 Identity 替代 fc 层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "563dddf8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 2048])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input = torch.randn(1, 1, 256, 256)\n",
    "output = resnet(input)\n",
    "output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "f4413f0c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[3.2843e-17, 0.0000e+00, 0.0000e+00, 4.0794e-17, 1.0000e+00]],\n",
       "        grad_fn=<SoftmaxBackward0>),\n",
       " tensor([[0.8848]], grad_fn=<AddmmBackward0>))"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = RESNET50_FC(input_channels=1, act_num=5, use_softmax=True)\n",
    "output = net(input)\n",
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4e0fbd9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch201",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
